{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MfBg1C5NB3X0"
   },
   "source": [
    "# Distributed training with TensorFlow\n",
    "\n",
    "**Learning Objectives**\n",
    "  1. Create MirroredStrategy\n",
    "  2. Integrate tf.distribute.Strategy with tf.keras\n",
    "  3. Create the input dataset and call tf.distribute.Strategy.experimental_distribute_dataset\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xHxb-dlhMIzW"
   },
   "source": [
    "## Introduction\n",
    "\n",
    "`tf.distribute.Strategy` is a TensorFlow API to distribute training \n",
    "across multiple GPUs, multiple machines or TPUs. Using this API, you can distribute your existing models and training code with minimal code changes.\n",
    "\n",
    "`tf.distribute.Strategy` has been designed with these key goals in mind:\n",
    "\n",
    "* Easy to use and support multiple user segments, including researchers, ML engineers, etc.\n",
    "* Provide good performance out of the box.\n",
    "* Easy switching between strategies.\n",
    "\n",
    "`tf.distribute.Strategy` can be used with a high-level API like [Keras](https://www.tensorflow.org/guide/keras), and can also be used to distribute custom training loops (and, in general, any computation using TensorFlow).\n",
    "\n",
    "In TensorFlow 2.x, you can execute your programs eagerly, or in a graph using [`tf.function`](function.ipynb). `tf.distribute.Strategy` intends to support both these modes of execution, but works best with `tf.function`. Eager mode is only recommended for debugging purpose and not supported for `TPUStrategy`. Although training is the focus of this guide, this API can also be used for distributing evaluation and prediction on different platforms.\n",
    "\n",
    "You can use `tf.distribute.Strategy` with very few changes to your code, because we have changed the underlying components of TensorFlow to become strategy-aware. This includes variables, layers, models, optimizers, metrics, summaries, and checkpoints.\n",
    "\n",
    "In this guide, we explain various types of strategies and how you can use them in different situations. To learn how to debug performance issues, see the [Optimize TensorFlow GPU Performance](gpu_performance_analysis.md) guide.\n",
    "\n",
    "Note: For a deeper understanding of the concepts, please watch [this deep-dive presentation](https://youtu.be/jKV53r9-H14). This is especially recommended if you plan to write your own training loop.\n",
    "\n",
    "Each learning objective will correspond to a __#TODO__ in the notebook where you will complete the notebook cell's code before running. Refer to the [solution](https://github.com/GoogleCloudPlatform/training-data-analyst/blob/master/courses/machine_learning/deepdive2/production_ml/solutions/distributed_training_with_TF.ipynb) for reference.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-02-13T02:22:53.263518Z",
     "iopub.status.busy": "2021-02-13T02:22:53.262758Z",
     "iopub.status.idle": "2021-02-13T02:23:00.382113Z",
     "shell.execute_reply": "2021-02-13T02:23:00.382626Z"
    },
    "id": "EVOZFbNgXghB"
   },
   "outputs": [],
   "source": [
    "# Import TensorFlow\n",
    "import tensorflow as tf"
   ]
  },
  {
 "cell_type": "markdown",
 "metadata": {
  "id": "m7KBpffWzlxH"
 },
 "source": [
  "This notebook uses TF2.x. Please check your tensorflow version using the cell below."
 ]
},
{
 "cell_type": "code",
 "execution_count": null,
 "metadata": {},
 "outputs": [
   {
    "name": "stdout",
    "output_type": "stream",
    "text": [
     "2.6.0\n"
    ]
   }
 ],
 "source": [
  "# Show the currently installed version of TensorFlow\n",
  "print(tf.__version__)"
 ]
},
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eQ1QESxxEbCh"
   },
   "source": [
    "## Types of strategies\n",
    "`tf.distribute.Strategy` intends to cover a number of use cases along different axes. Some of these combinations are currently supported and others will be added in the future. Some of these axes are:\n",
    "\n",
    "* *Synchronous vs asynchronous training:* These are two common ways of distributing training with data parallelism. In sync training, all workers train over different slices of input data in sync, and aggregating gradients at each step. In async training, all workers are independently training over the input data and updating variables asynchronously. Typically sync training is supported via all-reduce and async through parameter server architecture.\n",
    "* *Hardware platform:* You may want to scale your training onto multiple GPUs on one machine, or multiple machines in a network (with 0 or more GPUs each), or on Cloud TPUs.\n",
    "\n",
    "In order to support these use cases, there are six strategies available. The next section explains which of these are supported in which scenarios in TF. Here is a quick overview:\n",
    "\n",
    "| Training API          \t| MirroredStrategy  \t| TPUStrategy         \t| MultiWorkerMirroredStrategy     \t| CentralStorageStrategy          \t| ParameterServerStrategy  \t|\n",
    "|:-----------------------\t|:-------------------\t|:---------------------\t|:---------------------------------\t|:---------------------------------\t|:--------------------------\t|\n",
    "| **Keras API**             \t| Supported         \t|  Supported \t| Supported \t| Experimental support \t| Supported planned post 2.4 \t|\n",
    "| **Custom training loop**  \t| Supported \t| Supported   \t| Supported            \t| Experimental support           \t| Experimental support         \t|\n",
    "| **Estimator API**         \t| Limited Support         \t| Not supported           \t| Limited Support                       \t| Limited Support                       \t| Limited Support                \t|\n",
    "\n",
    "Note: [Experimental support](https://www.tensorflow.org/guide/versions#what_is_not_covered) means the APIs are not covered by any compatibilities guarantees.\n",
    "\n",
    "Note: Estimator support is limited. Basic training and evaluation are experimental, and advanced features—such as scaffold—are not implemented. We recommend using Keras or custom training loops if a use case is not covered."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DoQKKK8dtfg6"
   },
   "source": [
    "### MirroredStrategy\n",
    "`tf.distribute.MirroredStrategy` supports synchronous distributed training on multiple GPUs on one machine. It creates one replica per GPU device. Each variable in the model is mirrored across all the replicas. Together, these variables form a single conceptual variable called `MirroredVariable`. These variables are kept in sync with each other by applying identical updates.\n",
    "\n",
    "Efficient all-reduce algorithms are used to communicate the variable updates across the devices.\n",
    "All-reduce aggregates tensors across all the devices by adding them up, and makes them available on each device.\n",
    "It’s a fused algorithm that is very efficient and can reduce the overhead of synchronization significantly. There are many all-reduce algorithms and implementations available, depending on the type of communication available between devices. By default, it uses NVIDIA NCCL as the all-reduce implementation. You can choose from a few other options, or write your own.\n",
    "\n",
    "Here is the simplest way of creating `MirroredStrategy`:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-02-13T02:23:09.589761Z",
     "iopub.status.busy": "2021-02-13T02:23:09.588934Z",
     "iopub.status.idle": "2021-02-13T02:23:09.597397Z",
     "shell.execute_reply": "2021-02-13T02:23:09.597864Z"
    },
    "id": "9Z4FMAY9ADxK"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)\n"
     ]
    }
   ],
   "source": [
    "mirrored_strategy = tf.distribute.MirroredStrategy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wldY4aFCAH4r"
   },
   "source": [
    "This will create a `MirroredStrategy` instance which will use all the GPUs that are visible to TensorFlow, and use NCCL as the cross device communication.\n",
    "\n",
    "If you wish to use only some of the GPUs on your machine, you can do so like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-02-13T02:23:09.604157Z",
     "iopub.status.busy": "2021-02-13T02:23:09.603492Z",
     "iopub.status.idle": "2021-02-13T02:23:09.608343Z",
     "shell.execute_reply": "2021-02-13T02:23:09.607847Z"
    },
    "id": "nbGleskCACv_"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Some requested devices in `tf.distribute.Strategy` are not visible to TensorFlow: /job:localhost/replica:0/task:0/device:GPU:0,/job:localhost/replica:0/task:0/device:GPU:1\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1')\n"
     ]
    }
   ],
   "source": [
    "mirrored_strategy = tf.distribute.MirroredStrategy(devices=[\"/gpu:0\", \"/gpu:1\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8-KDnrJLAhav"
   },
   "source": [
    "If you wish to override the cross device communication, you can do so using the `cross_device_ops` argument by supplying an instance of `tf.distribute.CrossDeviceOps`. Currently,  `tf.distribute.HierarchicalCopyAllReduce` and `tf.distribute.ReductionToOneDevice` are two options other than `tf.distribute.NcclAllReduce` which is the default."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-02-13T02:23:09.612634Z",
     "iopub.status.busy": "2021-02-13T02:23:09.611960Z",
     "iopub.status.idle": "2021-02-13T02:23:09.614917Z",
     "shell.execute_reply": "2021-02-13T02:23:09.614420Z"
    },
    "id": "6-xIOIpgBItn"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)\n"
     ]
    }
   ],
   "source": [
    "# TODO 1 - Your code goes here.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kPEBCMzsGaO5"
   },
   "source": [
    "### TPUStrategy\n",
    "`tf.distribute.TPUStrategy` lets you run your TensorFlow training on Tensor Processing Units (TPUs). TPUs are Google's specialized ASICs designed to dramatically accelerate machine learning workloads. They are available on Google Colab, the [TensorFlow Research Cloud](https://www.tensorflow.org/tfrc) and [Cloud TPU](https://cloud.google.com/tpu).\n",
    "\n",
    "In terms of distributed training architecture, `TPUStrategy` is the same `MirroredStrategy` - it implements synchronous distributed training. TPUs provide their own implementation of efficient all-reduce and other collective operations across multiple TPU cores, which are used in `TPUStrategy`.\n",
    "\n",
    "Here is how you would instantiate `TPUStrategy`:\n",
    "\n",
    "Note: To run this code in Colab, you should select TPU as the Colab runtime. See [TensorFlow TPU Guide](https://www.tensorflow.org/guide/tpu).\n",
    "\n",
    "```\n",
    "cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(\n",
    "    tpu=tpu_address)\n",
    "tf.config.experimental_connect_to_cluster(cluster_resolver)\n",
    "tf.tpu.experimental.initialize_tpu_system(cluster_resolver)\n",
    "tpu_strategy = tf.distribute.TPUStrategy(cluster_resolver)\n",
    "```\n",
    "\n",
    "\n",
    "The `TPUClusterResolver` instance helps locate the TPUs. In Colab, you don't need to specify any arguments to it.\n",
    "\n",
    "If you want to use this for Cloud TPUs:\n",
    "- You must specify the name of your TPU resource in the `tpu` argument.\n",
    "- You must initialize the tpu system explicitly at the *start* of the program. This is required before TPUs can be used for computation. Initializing the tpu system also wipes out the TPU memory, so it's important to complete this step first in order to avoid losing state."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8Xc3gyo0Bejd"
   },
   "source": [
    "### MultiWorkerMirroredStrategy\n",
    "\n",
    "`tf.distribute.MultiWorkerMirroredStrategy` is very similar to `MirroredStrategy`. It implements synchronous distributed training across multiple workers, each with potentially multiple GPUs. Similar to `tf.distribute.MirroredStrategy`, it creates copies of all variables in the model on each device across all workers.\n",
    "\n",
    "Here is the simplest way of creating `MultiWorkerMirroredStrategy`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-02-13T02:23:09.619112Z",
     "iopub.status.busy": "2021-02-13T02:23:09.618361Z",
     "iopub.status.idle": "2021-02-13T02:23:09.623182Z",
     "shell.execute_reply": "2021-02-13T02:23:09.622671Z"
    },
    "id": "m3a_6ebbEjre"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using MirroredStrategy with devices ('/device:GPU:0',)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Single-worker MultiWorkerMirroredStrategy with local_devices = ('/device:GPU:0',), communication = CommunicationImplementation.AUTO\n"
     ]
    }
   ],
   "source": [
    "strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bt94JBvhEr4s"
   },
   "source": [
    "`MultiWorkerMirroredStrategy` has two implementations for cross-device communications.  `CommunicationImplementation.RING` is RPC-based and supports both CPU and GPU.  `CommunicationImplementation.NCCL` uses [Nvidia's NCCL](https://developer.nvidia.com/nccl) and provides the state of art performance on GPU, but it doesn't support CPU.  `CollectiveCommunication.AUTO` defers the choice to Tensorflow.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0JiImlw3F77E"
   },
   "source": [
    "One of the key differences to get multi worker training going, as compared to multi-GPU training, is the multi-worker setup. The `TF_CONFIG` environment variable is the standard way in TensorFlow to specify the cluster configuration to each worker that is part of the cluster. Learn more about [setting up TF_CONFIG](#TF_CONFIG)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3ZLBhaP9NUNr"
   },
   "source": [
    "### ParameterServerStrategy\n",
    "Parameter server training is a common data-parallel method to scale up model training on multiple machines. A parameter server training cluster consists of workers and parameter servers. Variables are created on parameter servers and they are read and updated by workers in each step. Please see the [parameter server training tutorial](../tutorials/distribute/parameter_server_training.ipynb) for details.\n",
    "\n",
    "TensorFlow 2 parameter server training uses a central-coordinator based architecture via the `tf.distribute.experimental.coordinator.ClusterCoordinator` class. \n",
    "\n",
    "In this implementation the `worker` and `parameter server` tasks run `tf.distribute.Server`s that listen for tasks from the coordinator. The coordinator creates resources, dispatches training tasks, writes checkpoints, and deals with task failures.\n",
    "\n",
    "In the programming running on the coordinator, you will use a `ParameterServerStrategy` object to define a training step and use a `ClusterCoordinator` to dispatch training steps to remote workers. Here is the simplest way to create them:\n",
    "\n",
    "```Python\n",
    "strategy = tf.distribute.experimental.ParameterServerStrategy(\n",
    "    tf.distribute.cluster_resolver.TFConfigClusterResolver(),\n",
    "    variable_partitioner=variable_partitioner)\n",
    "coordinator = tf.distribute.experimental.coordinator.ClusterCoordinator(\n",
    "    strategy)\n",
    "```\n",
    "\n",
    "Note you will need to configure TF_CONFIG environment variable if you use `TFConfigClusterResolver`. It is similar to  [TF_CONFIG](#TF_CONFIG) in `MultiWorkerMirroredStrategy` but has additional caveats.\n",
    "\n",
    "In TF 1, `ParameterServerStrategy` is available only with estimator via `tf.compat.v1.distribute.experimental.ParameterServerStrategy` symbol."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "E20tG21LFfv1"
   },
   "source": [
    "Note: This strategy is [`experimental`](https://www.tensorflow.org/guide/versions#what_is_not_covered) as it is currently under active development."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "45H0Wa8WKI8z"
   },
   "source": [
    "### CentralStorageStrategy\n",
    "`tf.distribute.experimental.CentralStorageStrategy` does synchronous training as well. Variables are not mirrored, instead they are placed on the CPU and operations are replicated across all local GPUs. If there is only one GPU, all variables and operations will be placed on that GPU.\n",
    "\n",
    "Create an instance of `CentralStorageStrategy` by:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-02-13T02:23:09.636243Z",
     "iopub.status.busy": "2021-02-13T02:23:09.635592Z",
     "iopub.status.idle": "2021-02-13T02:23:09.638242Z",
     "shell.execute_reply": "2021-02-13T02:23:09.638658Z"
    },
    "id": "rtjZOyaoMWrP"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:ParameterServerStrategy (CentralStorageStrategy if you are using a single machine) with compute_devices = ['/job:localhost/replica:0/task:0/device:GPU:0'], variable_device = '/job:localhost/replica:0/task:0/device:GPU:0'\n"
     ]
    }
   ],
   "source": [
    "central_storage_strategy = tf.distribute.experimental.CentralStorageStrategy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KY1nJHNkMl7b"
   },
   "source": [
    "This will create a `CentralStorageStrategy` instance which will use all visible GPUs and CPU. Update to variables on replicas will be aggregated before being applied to variables."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aAFycYUiNCUb"
   },
   "source": [
    "Note: This strategy is [`experimental`](https://www.tensorflow.org/guide/versions#what_is_not_covered) as it is currently a work in progress. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "t2XUdmIxKljq"
   },
   "source": [
    "### Other strategies\n",
    "\n",
    "In addition to the above strategies, there are two other strategies which might be useful for prototyping and debugging when using `tf.distribute` APIs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UD5I1beTpc7a"
   },
   "source": [
    "#### Default Strategy\n",
    "\n",
    "Default strategy is a distribution strategy which is present when no explicit distribution strategy is in scope. It implements the `tf.distribute.Strategy` interface but is a pass-through and provides no actual distribution. For instance, `strategy.run(fn)` will simply call `fn`. Code written using this strategy should behave exactly as code written without any strategy. You can think of it as a \"no-op\" strategy.\n",
    "\n",
    "Default strategy is a singleton - and one cannot create more instances of it. It can be obtained using `tf.distribute.get_strategy()` outside any explicit strategy's scope (the same API that can be used to get the current strategy inside an explicit strategy's scope). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-02-13T02:23:09.643033Z",
     "iopub.status.busy": "2021-02-13T02:23:09.642350Z",
     "iopub.status.idle": "2021-02-13T02:23:09.645348Z",
     "shell.execute_reply": "2021-02-13T02:23:09.644759Z"
    },
    "id": "ibHleFOOmPn9"
   },
   "outputs": [],
   "source": [
    "default_strategy = tf.distribute.get_strategy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EkxPl_5ImLzc"
   },
   "source": [
    "This strategy serves two main purposes:\n",
    "\n",
    "* It allows writing distribution aware library code unconditionally. For example, in `tf.optimizer`s can use `tf.distribute.get_strategy()` and use that strategy for reducing gradients - it will always return a strategy object on which we can call the reduce API.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-02-13T02:23:09.655662Z",
     "iopub.status.busy": "2021-02-13T02:23:09.654900Z",
     "iopub.status.idle": "2021-02-13T02:23:09.658340Z",
     "shell.execute_reply": "2021-02-13T02:23:09.658785Z"
    },
    "id": "WECeRzUdT6bU"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.0"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# In optimizer or other library code\n",
    "# Get currently active strategy\n",
    "strategy = tf.distribute.get_strategy()\n",
    "strategy.reduce(\"SUM\", 1., axis=None)  # reduce some values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JURbH-pUT51B"
   },
   "source": [
    "* Similar to library code, it can be used to write end users' programs to work with and without distribution strategy, without requiring conditional logic. A sample code snippet illustrating this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-02-13T02:23:09.663701Z",
     "iopub.status.busy": "2021-02-13T02:23:09.662970Z",
     "iopub.status.idle": "2021-02-13T02:23:09.970696Z",
     "shell.execute_reply": "2021-02-13T02:23:09.970059Z"
    },
    "id": "O4Vmae5jmSE6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MirroredVariable:{\n",
      "  0: <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "if tf.config.list_physical_devices('GPU'):\n",
    "  strategy = tf.distribute.MirroredStrategy()\n",
    "else:  # use default strategy\n",
    "  strategy = tf.distribute.get_strategy() \n",
    "\n",
    "with strategy.scope():\n",
    "  # do something interesting\n",
    "  print(tf.Variable(1.))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kTzsqN4lmJ0d"
   },
   "source": [
    "#### OneDeviceStrategy\n",
    "`tf.distribute.OneDeviceStrategy` is a strategy to place all variables and computation on a single specified device. \n",
    "\n",
    "```\n",
    "strategy = tf.distribute.OneDeviceStrategy(device=\"/gpu:0\")\n",
    "```\n",
    "\n",
    "This strategy is distinct from the default strategy in a number of ways. In default strategy, the variable placement logic remains unchanged when compared to running TensorFlow without any distribution strategy. But when using `OneDeviceStrategy`, all variables created in its scope are explicitly placed on the specified device. Moreover, any functions called via `OneDeviceStrategy.run` will also be placed on the specified device. \n",
    "\n",
    "Input distributed through this strategy will be prefetched to the specified device. In default strategy, there is no input distribution.\n",
    "\n",
    "Similar to the default strategy, this strategy could also be used to test your code before switching to other strategies which actually distribute to multiple devices/machines. This will exercise the distribution strategy machinery somewhat more than default strategy, but not to the full extent as using `MirroredStrategy` or `TPUStrategy` etc. If you want code that behaves as if no strategy, then use default strategy."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hQv1lm9UPDFy"
   },
   "source": [
    "So far you've seen the different strategies available and how you can instantiate them. The next few sections show the different ways in which you can use them to distribute your training. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_mcuy3UhPcen"
   },
   "source": [
    "## Using `tf.distribute.Strategy` with `tf.keras.Model.fit`\n",
    "\n",
    "`tf.distribute.Strategy` is integrated into `tf.keras` which is TensorFlow's implementation of the\n",
    "[Keras API specification](https://keras.io). `tf.keras`  is a high-level API to build and train models. By integrating into `tf.keras` backend, we've made it seamless for you to distribute your training written in the Keras training framework using `model.fit`.\n",
    "\n",
    "Here's what you need to change in your code:\n",
    "\n",
    "1. Create an instance of the appropriate `tf.distribute.Strategy`.\n",
    "2. Move the creation of Keras model, optimizer and metrics inside `strategy.scope`.\n",
    "\n",
    "We support all types of Keras models - sequential, functional and subclassed.\n",
    "\n",
    "Here is a snippet of code to do this for a very simple Keras model with one dense layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-02-13T02:23:09.976538Z",
     "iopub.status.busy": "2021-02-13T02:23:09.975879Z",
     "iopub.status.idle": "2021-02-13T02:23:10.043024Z",
     "shell.execute_reply": "2021-02-13T02:23:10.043470Z"
    },
    "id": "gbbcpzRnPZ6V"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n"
     ]
    }
   ],
   "source": [
    "mirrored_strategy = tf.distribute.MirroredStrategy()\n",
    "\n",
    "# TODO 2 - Your code goes here.\n",
    "\n",
    "model.compile(loss='mse', optimizer='sgd')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "773EOxCRVlTg"
   },
   "source": [
    "This example usees `MirroredStrategy` so you can run this on a machine with multiple GPUs. `strategy.scope()` indicates to Keras which strategy to use to distribute the training. Creating models/optimizers/metrics inside this scope allows us to create distributed variables instead of regular variables. Once this is set up, you can fit your model like you would normally. `MirroredStrategy` takes care of replicating the model's training on the available GPUs, aggregating gradients, and more."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-02-13T02:23:10.051847Z",
     "iopub.status.busy": "2021-02-13T02:23:10.051190Z",
     "iopub.status.idle": "2021-02-13T02:23:14.246874Z",
     "shell.execute_reply": "2021-02-13T02:23:14.246334Z"
    },
    "id": "ZMmxEFRTEjH5"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r",
      " 1/10 [==>...........................] - ETA: 24s - loss: 2.7806"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\r",
      "10/10 [==============================] - 3s 2ms/step - loss: 2.3180\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2/2\n",
      "\r",
      " 1/10 [==>...........................] - ETA: 0s - loss: 1.2290"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\r",
      "10/10 [==============================] - 0s 2ms/step - loss: 1.0246\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r",
      " 1/10 [==>...........................] - ETA: 10s - loss: 0.5432"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\r",
      "10/10 [==============================] - 1s 2ms/step - loss: 0.5432\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0.5432368516921997"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(10)\n",
    "model.fit(dataset, epochs=2)\n",
    "model.evaluate(dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nofTLwyXWHK8"
   },
   "source": [
    "Here a `tf.data.Dataset` provides the training and eval input. You can also use numpy arrays:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-02-13T02:23:14.252641Z",
     "iopub.status.busy": "2021-02-13T02:23:14.251965Z",
     "iopub.status.idle": "2021-02-13T02:23:15.488494Z",
     "shell.execute_reply": "2021-02-13T02:23:15.488947Z"
    },
    "id": "Lqgd9SdxW5OW"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r",
      " 1/10 [==>...........................] - ETA: 9s - loss: 0.5432"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\r",
      "10/10 [==============================] - 1s 2ms/step - loss: 0.3866\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2/2\n",
      "\r",
      " 1/10 [==>...........................] - ETA: 0s - loss: 0.2401"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\r",
      "10/10 [==============================] - 0s 2ms/step - loss: 0.1709\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.keras.callbacks.History at 0x7f36bfa9fc18>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "inputs, targets = np.ones((100, 1)), np.ones((100, 1))\n",
    "model.fit(inputs, targets, epochs=2, batch_size=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IKqaj7QwX0Zb"
   },
   "source": [
    "In both cases (dataset or numpy), each batch of the given input is divided equally among the multiple replicas. For instance, if using `MirroredStrategy` with 2 GPUs, each batch of size 10 will get divided among the 2 GPUs, with each receiving 5 input examples in each step. Each epoch will then train faster as you add more GPUs. Typically, you would want to increase your batch size as you add more accelerators so as to make effective use of the extra computing power. You will also need to re-tune your learning rate, depending on the model. You can use `strategy.num_replicas_in_sync` to get the number of replicas."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-02-13T02:23:15.495977Z",
     "iopub.status.busy": "2021-02-13T02:23:15.494996Z",
     "iopub.status.idle": "2021-02-13T02:23:15.498709Z",
     "shell.execute_reply": "2021-02-13T02:23:15.498167Z"
    },
    "id": "quNNTytWdGBf"
   },
   "outputs": [],
   "source": [
    "# Compute global batch size using number of replicas.\n",
    "BATCH_SIZE_PER_REPLICA = 5\n",
    "global_batch_size = (BATCH_SIZE_PER_REPLICA *\n",
    "                     mirrored_strategy.num_replicas_in_sync)\n",
    "dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100)\n",
    "dataset = dataset.batch(global_batch_size)\n",
    "\n",
    "LEARNING_RATES_BY_BATCH_SIZE = {5: 0.1, 10: 0.15}\n",
    "learning_rate = LEARNING_RATES_BY_BATCH_SIZE[global_batch_size]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z1Muy0gDZwO5"
   },
   "source": [
    "### What's supported now?\n",
    "\n",
    "\n",
    "| Training API \t| MirroredStrategy  \t| TPUStrategy         \t| MultiWorkerMirroredStrategy     \t| ParameterServerStrategy          \t| CentralStorageStrategy \t|\n",
    "|----------------\t|---------------------\t|-----------------------\t|-----------------------------------\t|-----------------------------------\t|---------------------------\t|\n",
    "| Keras APIs   \t| Supported \t| Supported \t| Experimental support \t| Experimental support \t| Experimental support   \t|\n",
    "\n",
    "### Examples and Tutorials\n",
    "\n",
    "Here is a list of tutorials and examples that illustrate the above integration end to end with Keras:\n",
    "\n",
    "1. [Tutorial](https://www.tensorflow.org/tutorials/distribute/keras) to train MNIST with `MirroredStrategy`.\n",
    "2. [Tutorial](https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras) to train MNIST using `MultiWorkerMirroredStrategy`.\n",
    "3. [Guide](https://www.tensorflow.org/guide/tpu#train_a_model_using_keras_high_level_apis) on training MNIST using `TPUStrategy`.\n",
    "4. [Tutorial](https://www.tensorflow.org/tutorials/distribute/parameter_server_training) for parameter server training in TensorFlow 2 with `ParameterServerStrategy`.\n",
    "5. TensorFlow Model Garden [repository](https://github.com/tensorflow/models/tree/master/official) containing collections of state-of-the-art models implemented using various strategies."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IlYVC0goepdk"
   },
   "source": [
    "## Using `tf.distribute.Strategy` with custom training loops\n",
    "As you've seen, using `tf.distribute.Strategy` with Keras `model.fit` requires changing only a couple lines of your code. With a little more effort, you can also use `tf.distribute.Strategy` with custom training loops.\n",
    "\n",
    "If you need more flexibility and control over your training loops than is possible with Estimator or Keras, you can write custom training loops. For instance, when using a GAN, you may want to take a different number of generator or discriminator steps each round. Similarly, the high level frameworks are not very suitable for Reinforcement Learning training.\n",
    "\n",
    "The `tf.distribute.Strategy` classes provide a core set of methods through to support custom training loops. Using these may require minor restructuring of the code initially, but once that is done, you should be able to switch between GPUs, TPUs, and multiple machines simply by changing the strategy instance.\n",
    "\n",
    "Here we will show a brief snippet illustrating this use case for a simple training example using the same Keras model as before.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XNHvSY32nVBi"
   },
   "source": [
    "First, create the model and optimizer inside the strategy's scope. This ensures that any variables created with the model and optimizer are mirrored variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-02-13T02:23:15.506636Z",
     "iopub.status.busy": "2021-02-13T02:23:15.505863Z",
     "iopub.status.idle": "2021-02-13T02:23:15.519213Z",
     "shell.execute_reply": "2021-02-13T02:23:15.518667Z"
    },
    "id": "W-3Bn-CaiPKD"
   },
   "outputs": [],
   "source": [
    "with mirrored_strategy.scope():\n",
    "  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])\n",
    "  optimizer = tf.keras.optimizers.SGD()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mYkAyPeYnlXk"
   },
   "source": [
    "Next, create the input dataset and call `tf.distribute.Strategy.experimental_distribute_dataset` to distribute the dataset based on the strategy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-02-13T02:23:15.524256Z",
     "iopub.status.busy": "2021-02-13T02:23:15.523591Z",
     "iopub.status.idle": "2021-02-13T02:23:15.535992Z",
     "shell.execute_reply": "2021-02-13T02:23:15.535446Z"
    },
    "id": "94BkvkLInkKd"
   },
   "outputs": [],
   "source": [
    "# TODO 3 - Your code goes here.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "grzmTlSvn2j8"
   },
   "source": [
    "Then, define one step of the training. Use `tf.GradientTape` to compute gradients and optimizer to apply those gradients to update our model's variables. To distribute this training step, put it in a function `train_step` and pass it to `tf.distrbute.Strategy.run` along with the dataset inputs you got from the `dist_dataset` created before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-02-13T02:23:15.544679Z",
     "iopub.status.busy": "2021-02-13T02:23:15.543845Z",
     "iopub.status.idle": "2021-02-13T02:23:15.546392Z",
     "shell.execute_reply": "2021-02-13T02:23:15.545884Z"
    },
    "id": "NJxL5YrVniDe"
   },
   "outputs": [],
   "source": [
    "loss_object = tf.keras.losses.BinaryCrossentropy(\n",
    "  from_logits=True,\n",
    "  reduction=tf.keras.losses.Reduction.NONE)\n",
    "\n",
    "def compute_loss(labels, predictions):\n",
    "  per_example_loss = loss_object(labels, predictions)\n",
    "  return tf.nn.compute_average_loss(per_example_loss, global_batch_size=global_batch_size)\n",
    "\n",
    "def train_step(inputs):\n",
    "  features, labels = inputs\n",
    "\n",
    "  with tf.GradientTape() as tape:\n",
    "    predictions = model(features, training=True)\n",
    "    loss = compute_loss(labels, predictions)\n",
    "\n",
    "  gradients = tape.gradient(loss, model.trainable_variables)\n",
    "  optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
    "  return loss\n",
    "\n",
    "@tf.function\n",
    "def distributed_train_step(dist_inputs):\n",
    "  per_replica_losses = mirrored_strategy.run(train_step, args=(dist_inputs,))\n",
    "  return mirrored_strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses,\n",
    "                         axis=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yRL5u_NLoTvq"
   },
   "source": [
    "A few other things to note in the code above:\n",
    "\n",
    "1. It used `tf.nn.compute_average_loss` to compute the loss. `tf.nn.compute_average_loss` sums the per example loss and divide the sum by the global_batch_size. This is important because later after the gradients are calculated on each replica, they are aggregated across the replicas by **summing** them.\n",
    "2. It used the `tf.distribute.Strategy.reduce` API to aggregate the results returned by `tf.distribute.Strategy.run`. `tf.distribute.Strategy.run` returns results from each local replica in the strategy, and there are multiple ways to consume this result. You can `reduce` them to get an aggregated value. You can also do `tf.distribute.Strategy.experimental_local_results` to get the list of values contained in the result, one per local replica.\n",
    "3. When `apply_gradients` is called within a distribution strategy scope, its behavior is modified. Specifically, before applying gradients on each parallel instance during synchronous training, it performs a sum-over-all-replicas of the gradients.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "o9k_6-6vpQ-P"
   },
   "source": [
    "Finally, once you have defined the training step, we can iterate over `dist_dataset` and run the training in a loop:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-02-13T02:23:15.551761Z",
     "iopub.status.busy": "2021-02-13T02:23:15.550312Z",
     "iopub.status.idle": "2021-02-13T02:23:17.888299Z",
     "shell.execute_reply": "2021-02-13T02:23:17.888716Z"
    },
    "id": "Egq9eufToRf6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(0.7167484, shape=(), dtype=float32)\n",
      "tf.Tensor(0.71152556, shape=(), dtype=float32)\n",
      "tf.Tensor(0.70635474, shape=(), dtype=float32)\n",
      "tf.Tensor(0.7012355, shape=(), dtype=float32)\n",
      "tf.Tensor(0.6961673, shape=(), dtype=float32)\n",
      "tf.Tensor(0.6911497, shape=(), dtype=float32)\n",
      "tf.Tensor(0.6861821, shape=(), dtype=float32)\n",
      "tf.Tensor(0.6812641, shape=(), dtype=float32)\n",
      "tf.Tensor(0.6763952, shape=(), dtype=float32)\n",
      "tf.Tensor(0.6715747, shape=(), dtype=float32)\n",
      "tf.Tensor(0.66680235, shape=(), dtype=float32)\n",
      "tf.Tensor(0.6620776, shape=(), dtype=float32)\n",
      "tf.Tensor(0.6573999, shape=(), dtype=float32)\n",
      "tf.Tensor(0.65276873, shape=(), dtype=float32)\n",
      "tf.Tensor(0.6481838, shape=(), dtype=float32)\n",
      "tf.Tensor(0.6436445, shape=(), dtype=float32)\n",
      "tf.Tensor(0.6391503, shape=(), dtype=float32)\n",
      "tf.Tensor(0.63470083, shape=(), dtype=float32)\n",
      "tf.Tensor(0.6302957, shape=(), dtype=float32)\n",
      "tf.Tensor(0.6259341, shape=(), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "for dist_inputs in dist_dataset:\n",
    "  print(distributed_train_step(dist_inputs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jK8eQXF_q1Zs"
   },
   "source": [
    "In the example above, you iterated over the `dist_dataset` to provide input to your training. We also provide the  `tf.distribute.Strategy.make_experimental_numpy_dataset` to support numpy inputs. You can use this API to create a dataset before calling `tf.distribute.Strategy.experimental_distribute_dataset`.\n",
    "\n",
    "Another way of iterating over your data is to explicitly use iterators. You may want to do this when you want to run for a given number of steps as opposed to iterating over the entire dataset.\n",
    "The above iteration would now be modified to first create an iterator and then explicitly call `next` on it to get the input data.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-02-13T02:23:17.894380Z",
     "iopub.status.busy": "2021-02-13T02:23:17.893015Z",
     "iopub.status.idle": "2021-02-13T02:23:17.956973Z",
     "shell.execute_reply": "2021-02-13T02:23:17.956340Z"
    },
    "id": "e5BEvR0-LJAc"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(0.62161595, shape=(), dtype=float32)\n",
      "tf.Tensor(0.6173406, shape=(), dtype=float32)\n",
      "tf.Tensor(0.6131077, shape=(), dtype=float32)\n",
      "tf.Tensor(0.60891676, shape=(), dtype=float32)\n",
      "tf.Tensor(0.60476726, shape=(), dtype=float32)\n",
      "tf.Tensor(0.60065883, shape=(), dtype=float32)\n",
      "tf.Tensor(0.59659094, shape=(), dtype=float32)\n",
      "tf.Tensor(0.5925633, shape=(), dtype=float32)\n",
      "tf.Tensor(0.5885753, shape=(), dtype=float32)\n",
      "tf.Tensor(0.5846267, shape=(), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "iterator = iter(dist_dataset)\n",
    "for _ in range(10):\n",
    "  print(distributed_train_step(next(iterator)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vDJO8mnypqBA"
   },
   "source": [
    "This covers the simplest case of using `tf.distribute.Strategy` API to distribute custom training loops."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BZjNwCt1qBdw"
   },
   "source": [
    "### What's supported now?\n",
    "\n",
    "| Training API          \t| MirroredStrategy  \t| TPUStrategy       \t| MultiWorkerMirroredStrategy \t| ParameterServerStrategy \t| CentralStorageStrategy \t| \n",
    "|:-----------------------\t|:-------------------\t|:-------------------\t|:-----------------------------\t|:------------------------\t|:-------------------------\t|\n",
    "| Custom Training Loop  \t|  Supported \t| Supported \t| Experimental support   \t| Experimental support \t| Experimental support        \t|\n",
    "\n",
    "### Examples and Tutorials\n",
    "Here are some examples for using distribution strategy with custom training loops:\n",
    "\n",
    "1. [Tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training) to train MNIST using `MirroredStrategy`.\n",
    "2. [Guide](https://www.tensorflow.org/guide/tpu#train_a_model_using_custom_training_loop) on training MNIST using `TPUStrategy`.\n",
    "3. TensorFlow Model Garden [repository](https://github.com/tensorflow/models/tree/master/official) containing collections of state-of-the-art models implemented using various strategies.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Xk0JdsTHyUnE"
   },
   "source": [
    "## Other topics\n",
    "\n",
    "This section covers some topics that are relevant to multiple use cases."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cP6BUIBtudRk"
   },
   "source": [
    "<a name=\"TF_CONFIG\"></a>\n",
    "### Setting up TF\\_CONFIG environment variable\n",
    "\n",
    "For multi-worker training, as mentioned before, you need to set `TF_CONFIG` environment variable for each\n",
    "binary running in your cluster. The `TF_CONFIG` environment variable is a JSON string which specifies what\n",
    "tasks constitute a cluster, their addresses and each task's role in the cluster.  The\n",
    "[tensorflow/ecosystem](https://github.com/tensorflow/ecosystem) repo provides a Kubernetes template in which sets\n",
    "`TF_CONFIG` for your training tasks.\n",
    "\n",
    "There are two components of TF_CONFIG: cluster and task. cluster provides information about the training cluster, which is a dict consisting of different types of jobs such as worker. In multi-worker training, there is usually one worker that takes on a little more responsibility like saving checkpoint and writing summary file for TensorBoard in addition to what a regular worker does. Such worker is referred to as the 'chief' worker, and it is customary that the worker with index 0 is appointed as the chief worker (in fact this is how tf.distribute.Strategy is implemented). task on the other hand provides information of the current task. The first component cluster is the same for all workers, and the second component task is different on each worker and specifies the type and index of that worker.\n",
    "\n",
    "One example of `TF_CONFIG` is:\n",
    "```\n",
    "os.environ[\"TF_CONFIG\"] = json.dumps({\n",
    "    \"cluster\": {\n",
    "        \"worker\": [\"host1:port\", \"host2:port\", \"host3:port\"],\n",
    "        \"ps\": [\"host4:port\", \"host5:port\"]\n",
    "    },\n",
    "   \"task\": {\"type\": \"worker\", \"index\": 1}\n",
    "})\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fezd3aF8wj9r"
   },
   "source": [
    "This `TF_CONFIG` specifies that there are three workers and two ps tasks in the\n",
    "cluster along with their hosts and ports. The \"task\" part specifies that the\n",
    "role of the current task in the cluster, worker 1 (the second worker). Valid roles in a cluster is\n",
    "\"chief\", \"worker\", \"ps\" and \"evaluator\". There should be no \"ps\" job except when using `tf.distribute.experimental.ParameterServerStrategy`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GXIbqSW-sFVg"
   },
   "source": [
    "## What's next?\n",
    "\n",
    "`tf.distribute.Strategy` is actively under development. Try it out and provide and your feedback using [GitHub issues](https://github.com/tensorflow/tensorflow/issues/new)."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [
    "Tce3stUlHN0L"
   ],
   "name": "distributed_training.ipynb",
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
